package org.example.function;

import com.antgroup.geaflow.api.context.RuntimeContext;
import com.antgroup.geaflow.api.function.RichFunction;
import com.antgroup.geaflow.api.function.io.SinkFunction;
import org.example.constants.Constants;
import org.example.csv.OutCSVData;
import org.example.model.LongStringResult;
import org.example.util.CSVFaster;
import org.example.util.FileUtils;

import java.nio.file.Path;

public class CSVFileSinkV2<OUT extends Comparable<? super OUT>> extends RichFunction implements SinkFunction<OUT> {

    private Path path;

    private final String resultName;

    private OutCSVData outCSVData;

    public CSVFileSinkV2(String resultName) {
        this.resultName = resultName;
    }

    @Override
    public void open(RuntimeContext runtimeContext) {
        String outputDir = runtimeContext.getConfiguration().getString(Constants.OUTPUT_DIR);
        path = FileUtils.pathJoin(outputDir, this.resultName);
        outCSVData = new OutCSVData();
    }

    @Override
    public void close() {
        // long st = System.currentTimeMillis();
        // this.csvData.sortValues();
        this.outCSVData.sortLines();
        // long ed = System.currentTimeMillis();
        // System.out.println("sort time: " + (ed - st) + " ms");
        // CSVUtils.write(path, this.csvData);
        CSVFaster.writeV1(path, this.outCSVData);
        // System.out.println("write time: " + (System.currentTimeMillis() - ed) + " ms");
        // CSVFaster2.write(path, this.csvData);
    }

    @Override
    public void write(OUT out) throws Exception {
        this.outCSVData.addLine((LongStringResult) out);
    }
}
